{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bde83aef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "\n",
    "import os\n",
    "import json\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "import gradio as gr\n",
    "import sqlite3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c92a3f84",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)\n",
    "\n",
    "gemini_key = os.getenv(\"GOOGLE_API_KEY\")\n",
    "if gemini_key:\n",
    "    print(f\"GEMINI API Key exists\")\n",
    "else:\n",
    "    print(\"GEMINI API Key not set\")\n",
    "\n",
    "GEMINI_BASE_URL = \"https://generativelanguage.googleapis.com/v1beta/openai/\"\n",
    "MODEL = 'gemini-2.5-flash-preview-05-20'\n",
    "gemini = OpenAI(base_url=GEMINI_BASE_URL, api_key=gemini_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfc051e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initializing db\n",
    "DB = \"books.db\"\n",
    "\n",
    "with sqlite3.connect(DB) as conn:\n",
    "    cursor = conn.cursor()\n",
    "    cursor.execute('''\n",
    "    CREATE TABLE books (\n",
    "        title TEXT PRIMARY KEY,\n",
    "        author TEXT,\n",
    "        available BOOLEAN,\n",
    "        borrower TEXT\n",
    "    )\n",
    "    ''')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22984b0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_book(title, author):\n",
    "    with sqlite3.connect(DB) as conn:\n",
    "        cursor = conn.cursor()\n",
    "        try:\n",
    "            cursor.execute(\n",
    "                'INSERT INTO books (title, author, available) VALUES (?,?,?)',\n",
    "                (title, author, 1)\n",
    "            )\n",
    "            return f'Book \"{title}\" added successfully.'\n",
    "        except sqlite3.IntegrityError:\n",
    "            return f'Book \"{title}\" already exists.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "384801e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add some dummy books\n",
    "add_book(\"Python 101\", \"Michael Driscoll\")\n",
    "add_book(\"Clean Code\", \"Robert C. Martin\")\n",
    "add_book(\"Deep Learning 101\", \"Ian Goodfellow\")\n",
    "add_book(\"The Pragmatic Programmer\", \"Andrew Hunt\")\n",
    "add_book(\"Introduction to Algorithms\", \"Cormen, Leiserson, Rivest, Stein\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03d69a32",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_book_info(title):\n",
    "    with sqlite3.connect(DB) as conn:\n",
    "        cursor = conn.cursor()\n",
    "        cursor.execute(\n",
    "            'SELECT title, author, available, borrower FROM books WHERE title = ?',\n",
    "            (title,)\n",
    "        )\n",
    "        result = cursor.fetchone()\n",
    "        if not result:\n",
    "            return f'Book \"{title}\" not found.'\n",
    "\n",
    "        title, author, available, borrower = result\n",
    "        status = 'available' if available else f'borrowed by {borrower}'        \n",
    "        return f'\"{title}\" by {author} is {status}.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c2fed74",
   "metadata": {},
   "outputs": [],
   "source": [
    "def borrow_book(title, borrower_name):\n",
    "    with sqlite3.connect(DB) as conn:\n",
    "        cursor = conn.cursor()\n",
    "        cursor.execute('SELECT available FROM books WHERE title = ?', (title,))\n",
    "        result = cursor.fetchone()\n",
    "        if not result:\n",
    "            return f'Book \"{title}\" not found.'\n",
    "        if not result[0]:\n",
    "            return f'Sorry, \"{title}\" is already borrowed.'\n",
    "\n",
    "        cursor.execute(\n",
    "            'UPDATE books SET borrower=?, available=? WHERE title=?',\n",
    "            (borrower_name, 0, title)\n",
    "        )\n",
    "    return f'\"{title}\" has been borrowed by {borrower_name}.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "635e2561",
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_book(title):\n",
    "    with sqlite3.connect(DB) as conn:\n",
    "        cursor = conn.cursor()\n",
    "        cursor.execute('SELECT available FROM books WHERE title = ?', (title,))\n",
    "        result = cursor.fetchone()\n",
    "        if not result:\n",
    "            return f'Book \"{title}\" not found.'\n",
    "        if result[0]:\n",
    "            return f'\"{title}\" is already available.'\n",
    "\n",
    "        cursor.execute(\n",
    "            'UPDATE books SET borrower=?, available=? WHERE title=?',\n",
    "            (None, 1, title)\n",
    "        )\n",
    "    return f'\"{title}\" has been returned and is now available.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "349301aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "add_book_tool = {\n",
    "    \"name\": \"add_book\",\n",
    "    \"description\": \"Add a new book to the library database.\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"title\": {\"type\": \"string\", \"description\": \"The book's title\"},\n",
    "            \"author\": {\"type\": \"string\", \"description\": \"The author's name\"}\n",
    "        },\n",
    "        \"required\": [\"title\", \"author\"],\n",
    "        \"additionalProperties\": False\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dcee2f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_book_info_tool = {\n",
    "    \"name\": \"get_book_info\",\n",
    "    \"description\": \"Get information about a book (author, availability, borrower).\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"title\": {\"type\": \"string\", \"description\": \"The book's title\"}\n",
    "        },\n",
    "        \"required\": [\"title\"],\n",
    "        \"additionalProperties\": False\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cabc15b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "borrow_book_tool = {\n",
    "    \"name\": \"borrow_book\",\n",
    "    \"description\": \"Borrow a book from the library if it is available.\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"title\": {\"type\": \"string\", \"description\": \"The book's title\"},\n",
    "            \"borrower_name\": {\"type\": \"string\", \"description\": \"Name of the borrower\"}\n",
    "        },\n",
    "        \"required\": [\"title\", \"borrower_name\"],\n",
    "        \"additionalProperties\": False\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f66664cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "return_book_tool = {\n",
    "    \"name\": \"return_book\",\n",
    "    \"description\": \"Return a borrowed book to the library.\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"title\": {\"type\": \"string\", \"description\": \"The book's title\"}\n",
    "        },\n",
    "        \"required\": [\"title\"],\n",
    "        \"additionalProperties\": False\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62b32c73",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = \"\"\"\n",
    "You are a helpful library assistant. \n",
    "You can add books, provide book info, let users borrow books, and handle returns.\n",
    "Always respond in a friendly, clear way suitable for a chatbot interface.\n",
    "Use the database functions when appropriate to answer questions about books.\n",
    "\"\"\"\n",
    "\n",
    "tools = [\n",
    "    {\"type\": \"function\", \"function\": add_book_tool},\n",
    "    {\"type\": \"function\", \"function\": get_book_info_tool},\n",
    "    {\"type\": \"function\", \"function\": borrow_book_tool},\n",
    "    {\"type\": \"function\", \"function\": return_book_tool}\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2af1a9a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def handle_tool_calls(message):\n",
    "    responses = []\n",
    "    for tool_call in message.tool_calls:\n",
    "        arguments = json.loads(tool_call.function.arguments)\n",
    "        if tool_call.function.name == \"add_book\":\n",
    "            result = add_book(arguments.get(\"title\"), arguments.get(\"author\"))\n",
    "        elif tool_call.function.name == \"get_book_info\":\n",
    "            result = get_book_info(arguments.get(\"title\"))\n",
    "        elif tool_call.function.name == \"borrow_book\":\n",
    "            result = borrow_book(arguments.get(\"title\"), arguments.get(\"borrower_name\"))\n",
    "        elif tool_call.function.name == \"return_book\":\n",
    "            result = return_book(arguments.get(\"title\"))\n",
    "        else:\n",
    "            result = f\"Unknown tool: {tool_call.function.name}\"\n",
    "\n",
    "        responses.append({\n",
    "            \"role\": \"tool\",\n",
    "            \"content\": result,\n",
    "            \"tool_call_id\": tool_call.id\n",
    "        })\n",
    "    return responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81451fd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(message, history):\n",
    "    history = [{\"role\": h[\"role\"], \"content\": h[\"content\"]} for h in history]\n",
    "    messages = [{\"role\": \"system\", \"content\": system_message}] + history + [{\"role\": \"user\", \"content\": message}]\n",
    "\n",
    "    response = gemini.chat.completions.create(model=MODEL, messages=messages, tools=tools)\n",
    "\n",
    "    # handle any tool calls\n",
    "    while response.choices[0].finish_reason == \"tool_calls\":\n",
    "        message = response.choices[0].message\n",
    "        print(message)\n",
    "        responses = handle_tool_calls(message)\n",
    "        messages.append(message)\n",
    "        messages.extend(responses)\n",
    "        response = gemini.chat.completions.create(model=MODEL, messages=messages, tools=tools)\n",
    "\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a43f313",
   "metadata": {},
   "outputs": [],
   "source": [
    "gr.ChatInterface(fn=chat, type=\"messages\").launch()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
